package com.luncheng.mcp.core.factory;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.luncheng.mcp.core.McpClientServer;
import com.luncheng.mcp.core.enums.TransportTypeEnum;
import com.luncheng.mcp.core.model.McpServerConfig;
import com.luncheng.mcp.core.strategy.TransportStrategy;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpServerTransportProvider;

import java.text.SimpleDateFormat;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Consumer;

public class McpServerFactory {
    private static ObjectMapper mapper;

    static {
        mapper = new ObjectMapper();
        mapper.setDateFormat(new SimpleDateFormat("yyyy-MM-dd HH:mm:ss"));
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }

    public static McpSyncServer createServer(McpServerConfig mcpServerConfig, Consumer<WebMvcSseServerTransportProvider> consumer) {
        try {
            TransportStrategy strategy = TransportFactory.getTransport(mcpServerConfig.getTransportType());
            McpServerTransportProvider transport = strategy.createServerTransport(mapper, mcpServerConfig);
            McpSyncServer mcpSyncServer = McpServer.sync(transport)
                    .serverInfo(mcpServerConfig.getName(), "1.0.0") // 设置服务器标识和版本
                    .capabilities(McpSchema.ServerCapabilities.builder()
                            .resources(true, true) // 启用资源读写功能
                            .tools(true)          // 启用工具功能
                            .prompts(true)        // 启用提示功能
                            .logging()            // 启用日志功能
                            .build())
                    .build();
            addTool(mcpServerConfig.getClients(), mcpSyncServer);
            mcpSyncServer.loggingNotification(McpSchema.LoggingMessageNotification.builder()
                    .level(McpSchema.LoggingLevel.INFO)
                    .logger("custom-logger")
                    .data("Server initialized")
                    .build());
            if (TransportTypeEnum.SSE.equals(mcpServerConfig.getTransportType())) {
                consumer.accept((WebMvcSseServerTransportProvider) transport);
            }
            return mcpSyncServer;
        } catch (Exception e) {
            System.err.println("MCP服务器初始化失败: " + e.getMessage());
            e.printStackTrace();
            throw new RuntimeException("MCP服务器初始化失败", e);
        }
    }

    private static void addTool(List<String> clients, McpSyncServer mcpSyncServer) throws JsonProcessingException {
        Map<String, McpSyncClient> clientMap = McpClientServer.getClientMap();
        for (String client : clients) {
            McpSyncClient mcpSyncClient = clientMap.get(client);
            if (Objects.isNull(mcpSyncClient)) {
                continue;
            }
            McpSchema.ListToolsResult listToolsResult = mcpSyncClient.listTools();
            List<McpSchema.Tool> tools = listToolsResult.tools();
            for (McpSchema.Tool tool : tools) {
                McpServerFeatures.SyncToolSpecification syncToolSpecification = new McpServerFeatures.SyncToolSpecification(
                        new McpSchema.Tool(tool.name(), tool.description(), tool.inputSchema()),
                        (mcpSyncServerExchange, stringObjectMap) -> mcpSyncClient.callTool(new McpSchema.CallToolRequest(tool.name(), stringObjectMap)));
                mcpSyncServer.addTool(syncToolSpecification);
            }
        }
    }
}